package service

import (
	"bufio"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"

	goscanmgt "goscan-mgt"
	"goscan-mgt/model"
)

// ScanAsyncTask 用于异步执行rustscan并写入数据库
// 需要传入CIDR、ScanService、ServiceInfoService、ProjectID
// 调用Start()方法即可异步执行

type ScanAsyncTask struct {
	CIDR      string
	ProjectID uint
}

func NewScanAsyncTask(cidr string, projectID uint) *ScanAsyncTask {
	return &ScanAsyncTask{
		CIDR:      cidr,
		ProjectID: projectID,
	}
}

func getRustscanPath() (string, error) {
	tmpDir := os.TempDir()
	exeName := "rustscan"
	if runtime.GOOS == "windows" {
		exeName += ".exe"
	}
	tmpPath := filepath.Join(tmpDir, exeName)
	if _, err := os.Stat(tmpPath); os.IsNotExist(err) {
		err := os.WriteFile(tmpPath, goscanmgt.RustscanBin, 0755)
		if err != nil {
			return "", err
		}
	}
	log.Println("rustscanPath:", tmpPath)
	return tmpPath, nil
}

func (t *ScanAsyncTask) Start() {
	go func() {
		rustscanPath, err := getRustscanPath()
		if err != nil {
			log.Println("extractRustscan error:", err)
			return
		}
		cmd := exec.Command(rustscanPath, "-a", t.CIDR)
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Println("cmd.StdoutPipe() error:", err)
			return
		}
		if err := cmd.Start(); err != nil {
			log.Println("cmd.Start() error:", err)
			return
		}
		scanner := bufio.NewScanner(stdout)
		re := regexp.MustCompile(`(?i)(open|closed|filtered|unfiltered)\s+([\d.]+):(\d+)`)
		// 获取端口-服务映射
		serviceMap := make(map[int]string)

		serviceInfoService := NewServiceInfoService()
		infos, _ := serviceInfoService.GetAllServiceInfo()
		for _, info := range infos {
			serviceMap[info.Port] = info.Service
		}
		scanService := NewScanService()
		var i int = 0
		for scanner.Scan() {
			line := scanner.Text()
			log.Println("rustscan output:", line)
			matches := re.FindStringSubmatch(line)
			if len(matches) == 4 {
				portStatus := matches[1]
				ip := matches[2]
				portStr := matches[3]
				port, _ := strconv.Atoi(portStr)
				service, exists := serviceMap[port]
				if !exists {
					service = "unknown"
				}
				result := &model.ScanResult{
					IP:         ip,
					Port:       port,
					PortStatus: portStatus,
					Service:    service,
					ProjectID:  &t.ProjectID,
				}
				log.Println("result:", result)
				err = scanService.CreateScanResult(result)
				if err != nil {
					log.Println("scanService.CreateScanResult() error:", err)
				}
				i += 1
			}
		}
		projectService := NewProjectService()
		_ = projectService.UpdateProjectStatus(t.ProjectID, 1, i)
		log.Println("cmd.Wait()")
		_ = cmd.Wait()
	}()
}
